import cv2
import numpy as np
import math

class ContourProcessor:
    """
    一个集成了轮廓查找、选择、偏移和自适应细分功能的处理器类。

    该类将所有相关操作封装在一起，提供一个清晰的接口来处理二值图像，
    并生成一条经过偏移和自适应细分的路径。
    """

    def __init__(self, length_threshold=40.0, straight_step=50, curve_step=5):
        """
        初始化处理器。

        Args:
            length_threshold (float): 边长阈值，用于区分直线和曲线。
            straight_step (int): 直线段的采样步长。
            curve_step (int): 曲线段的采样步长。
        """
        # 细分参数
        self.length_threshold = length_threshold
        self.straight_step = max(1, straight_step)
        self.curve_step = max(1, curve_step)

        # 用于可视化的中间结果
        self.base_image = None
        self.selected_contour = None
        self.offset_contour = None
        self.approximated_poly = None
        self.subdivided_path = None

    def _find_and_select_contour(self, binary_image, contour_type, min_area=0, max_area=float('inf')):
        """
        在图像中查找并根据类型和面积筛选目标轮廓。
        """
        contours, hierarchy = cv2.findContours(binary_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return None

        pre_selected_contours = []
        if hierarchy is not None:
            for i, contour in enumerate(contours):
                is_outer = hierarchy[0][i][3] == -1
                if contour_type == 'outer' and is_outer:
                    pre_selected_contours.append(contour)
                elif contour_type == 'inner' and not is_outer:
                    pre_selected_contours.append(contour)
        
        if not pre_selected_contours:
            return None
        
        final_contours = []
        for c in pre_selected_contours:
            area = cv2.contourArea(c)
            if min_area <= area <= max_area:
                final_contours.append(c)

        if not final_contours:
            return None

        return max(final_contours, key=cv2.contourArea)

    def _offset_contour_morphological(self, contour, offset, image_shape):
        """
        使用形态学操作对轮廓进行内外偏移。
        """
        if offset == 0:
            return contour

        mask = np.zeros(image_shape, dtype=np.uint8)
        cv2.drawContours(mask, [contour], -1, 255, -1)
        
        kernel_size = abs(offset)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        
        offset_mask = cv2.dilate(mask, kernel, iterations=1) if offset > 0 else cv2.erode(mask, kernel, iterations=1)
            
        contours, _ = cv2.findContours(offset_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        return max(contours, key=cv2.contourArea) if contours else None

    def _offset_contour_by_scaling(self, contour, offset):
        """
        通过几何方法（点到中心的向量）对轮廓进行内外偏移。
        """
        if offset == 0:
            return contour

        M = cv2.moments(contour)
        if M["m00"] == 0:
            return contour 
        cx = int(M["m10"] / M["m00"])
        cy = int(M["m01"] / M["m00"])
        center = np.array([cx, cy])

        offset_points = []
        for point in contour:
            p = point[0]
            v = p - center
            dist = np.linalg.norm(v)
            if dist < 1e-6:
                offset_points.append(point)
                continue
            
            new_p = p + (v / dist) * offset
            offset_points.append([new_p.astype(int)])

        return np.array(offset_points)


    # def _reorder_contour_start_point(self, contour):
    #     """
    #     重新排序轮廓点，使其从最左上角的点开始。
    #     """
    #     if contour is None or len(contour) < 1:
    #         return contour
        
    #     reshaped_contour = contour.reshape(-1, 2)
    #     start_point_index = np.argmin(reshaped_contour.sum(axis=1))
        
    #     reordered_contour = np.roll(contour, -start_point_index, axis=0)
        
    #     return reordered_contour

    def _reorder_contour_start_point(self, contour):
        """
        使用更精确的字典序方法重新排序轮廓点，使其从最左上角的点开始。
        """
        if contour is None or len(contour) < 1:
            return contour
        
        # 将轮廓转换为 (N, 2) 的形状以便操作
        reshaped_contour = contour.reshape(-1, 2)
        
        # 1. 找到最小的 y 坐标
        min_y = np.min(reshaped_contour[:, 1])
        
        # 2. 筛选出所有 y 坐标最小的点
        top_points = reshaped_contour[reshaped_contour[:, 1] == min_y]
        
        # 3. 在这些点中，找到 x 坐标最小的点，这就是最左上角的点
        top_left_point = top_points[np.argmin(top_points[:, 0])]
        
        # 4. 找到这个点在原始轮廓数组中的索引
        # 使用 np.where 查找所有匹配的点，并取第一个
        start_point_index = np.where((reshaped_contour == top_left_point).all(axis=1))[0][0]
        
        # 5. “滚动”数组，使这个点成为新的起点
        reordered_contour = np.roll(contour, -start_point_index, axis=0)
        
        return reordered_contour

    def _get_adaptive_step(self, contour, index):
        """
        辅助函数：根据给定索引处的局部曲率返回合适的步长。
        """
        num_points = len(contour)
        p_prev = contour[(index - 5 + num_points) % num_points][0]
        p_curr = contour[index][0]
        p_next = contour[(index + 5) % num_points][0]
        
        v1 = p_curr - p_prev
        v2 = p_next - p_curr

        norm1 = np.linalg.norm(v1)
        norm2 = np.linalg.norm(v2)

        dot_product = 0.0
        if norm1 > 1e-6 and norm2 > 1e-6:
            dot_product = np.dot(v1 / norm1, v2 / norm2)

        is_straight = dot_product > 0.996
        return self.straight_step if is_straight else self.curve_step

    def _generate_final_path(self, contour, epsilon_factor, mode='simplified'):
        """
        根据指定模式生成最终路径。
        """
        if contour is None or len(contour) < 3:
            return np.array([])
        
        if mode == 'simplified':
            area = cv2.contourArea(contour)
            if area < 1e-6: return np.array([])
            
            epsilon = epsilon_factor * math.sqrt(area)
            self.approximated_poly = cv2.approxPolyDP(contour, epsilon, True)

            if not np.array_equal(self.approximated_poly[0], self.approximated_poly[-1]):
                 self.approximated_poly = np.vstack([self.approximated_poly, self.approximated_poly[0:1]])

            final_path = []
            for i in range(len(self.approximated_poly) - 1):
                p1 = self.approximated_poly[i][0]
                p2 = self.approximated_poly[i + 1][0]
                edge_vec = p2 - p1
                edge_length = np.linalg.norm(edge_vec)

                if edge_length < 1e-6: continue

                num_points = int(edge_length / self.straight_step) if edge_length > self.length_threshold else int(edge_length / self.curve_step)
                num_points = max(1, num_points)
                
                segment_points = np.linspace(p1, p2, num_points, endpoint=False)
                final_path.extend(segment_points)
            
            self.subdivided_path = np.array(final_path, dtype=np.int32)

        elif mode == 'direct':
            self.approximated_poly = None
            final_path = [contour[0][0]]
            dist_accumulator = 0.0
            
            closed_contour = np.vstack([contour, contour[0:1]])

            for i in range(len(closed_contour) - 1):
                p1 = closed_contour[i][0].astype(float)
                p2 = closed_contour[i+1][0].astype(float)
                
                segment_vec = p2 - p1
                segment_len = np.linalg.norm(segment_vec)
                
                if segment_len < 1e-6:
                    continue

                current_step = self._get_adaptive_step(contour, i)
                dist_accumulator += segment_len
                
                while dist_accumulator >= current_step:
                    dist_from_p1 = segment_len - (dist_accumulator - current_step)
                    
                    if segment_len > 1e-6:
                        ratio = dist_from_p1 / segment_len
                        new_point = p1 + ratio * segment_vec
                        final_path.append(new_point)
                    
                    dist_accumulator -= current_step
            
            self.subdivided_path = np.array(final_path, dtype=np.int32)
        
        else:
            print(f"未知的细分模式: {mode}. 返回空路径。")
            self.subdivided_path = np.array([])

        return self.subdivided_path


    def process_image(self, binary_image, offset_distance, contour_type='outer', min_area=30, max_area=float('inf'), epsilon_factor=0.04, offset_method='morphological', subdivision_mode='simplified', sort_direction='clockwise'):
        """
        处理二值图像的完整流程。
        binary_image: 输入的二值图像。
        offset_distance: 偏移距离，正值表示外偏移，负值表示内偏移。
        contour_type: 'outer' 或 'inner'，选择外轮廓或内轮廓。
        min_area: 最小面积筛选阈值。
        max_area: 最大面积筛选阈值。
        epsilon_factor: 用于简化轮廓的因子。
        offset_method: 'morphological' 或 'scaling'，选择偏移方法。
        subdivision_mode: 'simplified' 或 'direct'，选择细分模式。
        sort_direction: 'clockwise' 或 'counter-clockwise'，选择轮廓排序方向。
        """
        self.base_image = binary_image.copy()

        self.selected_contour = self._find_and_select_contour(binary_image, contour_type, min_area, max_area)
        if self.selected_contour is None:
            print(f"未找到符合条件的 '{contour_type}' 轮廓。")
            return []
        
        if offset_method == 'morphological':
            self.offset_contour = self._offset_contour_morphological(self.selected_contour, offset_distance, binary_image.shape[:2])
        elif offset_method == 'scaling':
            self.offset_contour = self._offset_contour_by_scaling(self.selected_contour, offset_distance)
        else:
            print(f"未知的偏移方法: {offset_method}. 将使用 'morphological'.")
            self.offset_contour = self._offset_contour_morphological(self.selected_contour, offset_distance, binary_image.shape[:2])

        if self.offset_contour is None:
            print("偏移后轮廓消失（可能偏移量过大）。")
            return []
        
        
        if sort_direction in ['clockwise', 'counter-clockwise']:
            # OpenCV中，面积为正表示顺时针(CW)，为负表示逆时针(CCW)
            area = cv2.contourArea(self.offset_contour, oriented=True)
            is_ccw = area < 0  # 负值说明是逆时针

            if sort_direction == 'clockwise' and is_ccw:
                # 期望顺时针，但当前是逆时针，则反转
                self.offset_contour = self.offset_contour[::-1]
            elif sort_direction == 'counter-clockwise' and not is_ccw:
                # 期望逆时针，但当前是顺时针，则反转
                self.offset_contour = self.offset_contour[::-1]

        self.offset_contour = self._reorder_contour_start_point(self.offset_contour)

        path_points = self._generate_final_path(self.offset_contour, epsilon_factor, mode=subdivision_mode)
        
        return list(enumerate(path_points))

    def visualize(self,image):
        """
        在原始输入图像上创建可视化结果。
        """
        if image is None:
            print("错误: 请先调用 process_image 方法再进行可视化。")
            return None

        if len(image.shape) == 2 or image.shape[2] == 1:
            vis_img = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        else:
            vis_img = image.copy()

        if self.selected_contour is not None:
            cv2.drawContours(vis_img, [self.selected_contour], -1, (180, 180, 180), 3)
        
        if self.offset_contour is not None:
            cv2.drawContours(vis_img, [self.offset_contour], -1, (0, 255, 0), 2)

        if self.approximated_poly is not None:
            cv2.polylines(vis_img, [self.approximated_poly], isClosed=True, color=(255, 0, 0), thickness=1)

        if self.subdivided_path is not None:
            for i, pt in enumerate(self.subdivided_path):
                cv2.circle(vis_img, tuple(pt), 3, (0, 0, 255), -1)
                if i < 5: # 为前5个点绘制序号以显示方向
                    cv2.putText(vis_img, str(i), (pt[0] + 5, pt[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)

        return vis_img


class PathFollower:
    def __init__(self, path, threshold=10):
        """
        path: list of (index, np.array([x, y])) 目标路径
        threshold: 判断接近目标点的距离阈值
        """
        self.path = path
        self.threshold = threshold
        self.current_index = 0
        self.first_round = False
        self.finished = False
        self.path_length = len(path)

    def reset(self):
        self.current_index = 0
        self.first_round = False
        self.finished = False

    def update(self, current_position):
        """
        current_position: tuple or np.array([x, y])
        返回值:
            - 当前目标点索引
            - 当前目标坐标
            - 是否到达目标点
            - 是否路径完成（绕了一圈）
        """
        if self.finished:
            return None, None, False, True
        
        if not self.path or self.current_index >= self.path_length:
            print("[FollowLine] 路径未初始化或索引超界")
            return None, None, False, True

        # 当前目标
        target_idx, target_coord = self.path[self.current_index]
        current_x, current_y = current_position
        target_x, target_y = target_coord

        # 计算距离
        dist = np.linalg.norm(np.array([current_x, current_y]) - target_coord)

        reached = dist < self.threshold
        completed = False

        if reached:
            print(f"到达目标点 {target_idx}: {target_coord.tolist()}")

            if self.current_index == 0:
                if self.first_round:
                    self.finished = True
                    completed = True
                else:
                    self.first_round = True

            # 前进到下一个目标点
            self.current_index = (self.current_index + 1) % self.path_length

        return target_idx, target_coord, reached, completed

    
# # # --- 使用示例 ---
# if __name__ == '__main__':

#     # 1. 读取二值图像
#     input_image = cv2.imread('/home/sunrise/cricket-training-questions/images/512.jpg')
#     # input_image = cv2.imread('/home/sunrise/cricket-training-questions/images/123.webp')
#     # input_image = cv2.imread('/home/sunrise/cricket-training-questions/images/1234.webp')
#     input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY)
#     _, input_image = cv2.threshold(input_image, 127, 255, cv2.THRESH_BINARY)
#     input_image = cv2.bitwise_not(input_image)  # 反转颜色，使目标为白色
#     cv2.imshow("Input Binary Image", input_image)
#     cv2.waitKey(0)
#     height, width = input_image.shape

#     processor = ContourProcessor(length_threshold=100, straight_step=30, curve_step=8)
    
#     # --- 场景1: 使用 'simplified' 模式 (默认) ---
#     print("--- 场景1: 使用 'simplified' 模式进行细分 ---")
#     indexed_path_simplified = processor.process_image(
#         input_image, 
#         offset_distance=-15,
#         contour_type='outer', 
#         min_area=1000,
#         epsilon_factor=0.02, # 因子
#         offset_method = "morphological", # 使用几何方法进行偏移
#         subdivision_mode='simplified', # 明确指定
#         sort_direction='counter-clockwise' # 顺时针排序
#     )
#     print(f"处理结果: {indexed_path_simplified}")
#     visualization_simplified = processor.visualize()
    
#     # --- 场景2: 使用 'direct' 模式 ---
#     print("\n--- 场景2: 使用 'direct' 模式进行细分 ---")
#     indexed_path_direct = processor.process_image(
#         input_image, 
#         offset_distance=-10,
#         contour_type='outer',
#         min_area=1000,
#         epsilon_factor=0.02, # 因子
#         offset_method = "scaling", # 使用几何方法进行偏移
#         subdivision_mode='direct' # <-- 使用新的直接细分模式
#     )
#     visualization_direct = processor.visualize()

#     # 可视化结果
#     if visualization_simplified is not None and visualization_direct is not None:
#         cv2.putText(visualization_simplified, "'Simplified' Mode", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
#         cv2.putText(visualization_direct, "'Direct' Mode (Adaptive)", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
        
#         # 将两个结果并排显示以便比较
#         combined_vis = np.hstack((visualization_simplified, visualization_direct))
#         cv2.imshow("Path Processing Comparison", combined_vis)
#         cv2.waitKey(0)
#         cv2.destroyAllWindows()
